{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/cookbooks/contextual_retrieval.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Contextual Retrieval\n",
    "\n",
    "In this notebook we will demonstrate how you can implement [Anthropic's Contextual Retrieval](https://www.anthropic.com/news/contextual-retrieval) using LlamaIndex abstractions.\n",
    "\n",
    "We will use:\n",
    "\n",
    "1. `Paul Graham Essay` dataset.\n",
    "2. Anthropic LLM for context creation for each chunk.\n",
    "3. OpenAI LLM for Synthetic query generation and embedding model.\n",
    "4. CohereAI Reranker."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Installation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -U llama-index llama-index-llms-anthropic llama-index-postprocessor-cohere-rerank llama-index-retrievers-bm25 stemmer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import nest_asyncio\n",
    "\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup API Keys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# For creating context for each chunk\n",
    "os.environ[\"ANTHROPIC_API_KEY\"] = \"<YOUR ANTHROPIC API KEY>\"\n",
    "\n",
    "# For creating synthetic dataset and embedding model\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"<YOUR OPENAI API KEY>\"\n",
    "\n",
    "# For reranker\n",
    "os.environ[\"COHERE_API_KEY\"] = \"<YOUR COHEREAI API KEY>\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup LLM and Embedding model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.anthropic import Anthropic\n",
    "\n",
    "llm_anthropic = Anthropic(model=\"claude-3-5-sonnet-20240620\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.embeddings.openai import OpenAIEmbedding\n",
    "from llama_index.core import Settings\n",
    "\n",
    "Settings.embed_model = OpenAIEmbedding()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2024-10-01 13:00:06--  https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt\n",
      "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.108.133, 185.199.109.133, 185.199.110.133, ...\n",
      "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.108.133|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 75042 (73K) [text/plain]\n",
      "Saving to: ‘./paul_graham_essay.txt’\n",
      "\n",
      "./paul_graham_essay 100%[===================>]  73.28K  --.-KB/s    in 0.08s   \n",
      "\n",
      "2024-10-01 13:00:06 (921 KB/s) - ‘./paul_graham_essay.txt’ saved [75042/75042]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt' -O './paul_graham_essay.txt'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import SimpleDirectoryReader\n",
    "\n",
    "documents = SimpleDirectoryReader(\n",
    "    input_files=[\"./paul_graham_essay.txt\"],\n",
    ").load_data()\n",
    "\n",
    "WHOLE_DOCUMENT = documents[0].text"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prompts for creating context for each chunk\n",
    "\n",
    "We will utilize anthropic prompt caching for creating context for each chunk. If you haven’t explored our integration yet, please take a moment to review it [here](https://github.com/run-llama/llama_index/blob/main/docs/examples/llm/anthropic_prompt_caching.ipynb)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt_document = \"\"\"<document>\n",
    "{WHOLE_DOCUMENT}\n",
    "</document>\"\"\"\n",
    "\n",
    "prompt_chunk = \"\"\"Here is the chunk we want to situate within the whole document\n",
    "<chunk>\n",
    "{CHUNK_CONTENT}\n",
    "</chunk>\n",
    "Please give a short succinct context to situate this chunk within the overall document for the purposes of improving search retrieval of the chunk. Answer only with the succinct context and nothing else.\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Utils\n",
    "\n",
    "1. `create_contextual_nodes` - Function to create contextual nodes for a list of nodes.\n",
    "\n",
    "2. `create_embedding_retriever` - Function to create an embedding retriever for a list of nodes.\n",
    "\n",
    "3. `create_bm25_retriever` - Function to create a bm25 retriever for a list of nodes.\n",
    "\n",
    "4. `EmbeddingBM25RerankerRetriever` - Custom retriever that uses both embedding and bm25 retrievers and reranker.\n",
    "\n",
    "5. `create_eval_dataset` - Function to create a evaluation dataset from a list of nodes.\n",
    "\n",
    "6. `set_node_ids` - Function to set node ids for a list of nodes.\n",
    "\n",
    "7. `retrieval_results` - Function to get retrieval results for a retriever and evaluation dataset.\n",
    "\n",
    "8. `display_results` - Function to display results from `retrieval_results`\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.retrievers.bm25 import BM25Retriever\n",
    "from llama_index.core.evaluation import (\n",
    "    generate_question_context_pairs,\n",
    "    RetrieverEvaluator,\n",
    ")\n",
    "from llama_index.core.retrievers import BaseRetriever, VectorIndexRetriever\n",
    "from llama_index.core.schema import NodeWithScore\n",
    "from llama_index.core import VectorStoreIndex, QueryBundle\n",
    "from llama_index.core.llms import ChatMessage, TextBlock\n",
    "\n",
    "import pandas as pd\n",
    "import copy\n",
    "import Stemmer\n",
    "\n",
    "from typing import List\n",
    "\n",
    "\n",
    "def create_contextual_nodes(nodes_):\n",
    "    \"\"\"Function to create contextual nodes for a list of nodes\"\"\"\n",
    "    nodes_modified = []\n",
    "    for node in nodes_:\n",
    "        new_node = copy.deepcopy(node)\n",
    "        messages = [\n",
    "            ChatMessage(role=\"system\", content=\"You are helpful AI Assitant.\"),\n",
    "            ChatMessage(\n",
    "                role=\"user\",\n",
    "                content=[\n",
    "                    TextBlock(\n",
    "                        text=prompt_document.format(\n",
    "                            WHOLE_DOCUMENT=WHOLE_DOCUMENT\n",
    "                        )\n",
    "                    ),\n",
    "                    TextBlock(\n",
    "                        text=prompt_chunk.format(CHUNK_CONTENT=node.text)\n",
    "                    ),\n",
    "                ],\n",
    "                additional_kwargs={\"cache_control\": {\"type\": \"ephemeral\"}},\n",
    "            ),\n",
    "        ]\n",
    "        new_node.metadata[\"context\"] = str(\n",
    "            llm_anthropic.chat(\n",
    "                messages,\n",
    "                extra_headers={\"anthropic-beta\": \"prompt-caching-2024-07-31\"},\n",
    "            )\n",
    "        )\n",
    "        nodes_modified.append(new_node)\n",
    "\n",
    "    return nodes_modified\n",
    "\n",
    "\n",
    "def create_embedding_retriever(nodes_, similarity_top_k=2):\n",
    "    \"\"\"Function to create an embedding retriever for a list of nodes\"\"\"\n",
    "    vector_index = VectorStoreIndex(nodes_)\n",
    "    retriever = vector_index.as_retriever(similarity_top_k=similarity_top_k)\n",
    "    return retriever\n",
    "\n",
    "\n",
    "def create_bm25_retriever(nodes_, similarity_top_k=2):\n",
    "    \"\"\"Function to create a bm25 retriever for a list of nodes\"\"\"\n",
    "    bm25_retriever = BM25Retriever.from_defaults(\n",
    "        nodes=nodes_,\n",
    "        similarity_top_k=similarity_top_k,\n",
    "        stemmer=Stemmer.Stemmer(\"english\"),\n",
    "        language=\"english\",\n",
    "    )\n",
    "    return bm25_retriever\n",
    "\n",
    "\n",
    "def create_eval_dataset(nodes_, llm, num_questions_per_chunk=2):\n",
    "    \"\"\"Function to create a evaluation dataset for a list of nodes\"\"\"\n",
    "    qa_dataset = generate_question_context_pairs(\n",
    "        nodes_, llm=llm, num_questions_per_chunk=num_questions_per_chunk\n",
    "    )\n",
    "    return qa_dataset\n",
    "\n",
    "\n",
    "def set_node_ids(nodes_):\n",
    "    \"\"\"Function to set node ids for a list of nodes\"\"\"\n",
    "\n",
    "    # by default, the node ids are set to random uuids. To ensure same id's per run, we manually set them.\n",
    "    for index, node in enumerate(nodes_):\n",
    "        node.id_ = f\"node_{index}\"\n",
    "\n",
    "    return nodes_\n",
    "\n",
    "\n",
    "async def retrieval_results(retriever, eval_dataset):\n",
    "    \"\"\"Function to get retrieval results for a retriever and evaluation dataset\"\"\"\n",
    "\n",
    "    metrics = [\"hit_rate\", \"mrr\", \"precision\", \"recall\", \"ap\", \"ndcg\"]\n",
    "\n",
    "    retriever_evaluator = RetrieverEvaluator.from_metric_names(\n",
    "        metrics, retriever=retriever\n",
    "    )\n",
    "\n",
    "    eval_results = await retriever_evaluator.aevaluate_dataset(qa_dataset)\n",
    "\n",
    "    return eval_results\n",
    "\n",
    "\n",
    "def display_results(name, eval_results):\n",
    "    \"\"\"Display results from evaluate.\"\"\"\n",
    "\n",
    "    metrics = [\"hit_rate\", \"mrr\", \"precision\", \"recall\", \"ap\", \"ndcg\"]\n",
    "\n",
    "    metric_dicts = []\n",
    "    for eval_result in eval_results:\n",
    "        metric_dict = eval_result.metric_vals_dict\n",
    "        metric_dicts.append(metric_dict)\n",
    "\n",
    "    full_df = pd.DataFrame(metric_dicts)\n",
    "\n",
    "    columns = {\n",
    "        \"retrievers\": [name],\n",
    "        **{k: [full_df[k].mean()] for k in metrics},\n",
    "    }\n",
    "\n",
    "    metric_df = pd.DataFrame(columns)\n",
    "\n",
    "    return metric_df\n",
    "\n",
    "\n",
    "class EmbeddingBM25RerankerRetriever(BaseRetriever):\n",
    "    \"\"\"Custom retriever that uses both embedding and bm25 retrievers and reranker\"\"\"\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        vector_retriever: VectorIndexRetriever,\n",
    "        bm25_retriever: BM25Retriever,\n",
    "        reranker: CohereRerank,\n",
    "    ) -> None:\n",
    "        \"\"\"Init params.\"\"\"\n",
    "\n",
    "        self._vector_retriever = vector_retriever\n",
    "        self.bm25_retriever = bm25_retriever\n",
    "        self.reranker = reranker\n",
    "\n",
    "        super().__init__()\n",
    "\n",
    "    def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:\n",
    "        \"\"\"Retrieve nodes given query.\"\"\"\n",
    "\n",
    "        vector_nodes = self._vector_retriever.retrieve(query_bundle)\n",
    "        bm25_nodes = self.bm25_retriever.retrieve(query_bundle)\n",
    "\n",
    "        vector_nodes.extend(bm25_nodes)\n",
    "\n",
    "        retrieved_nodes = self.reranker.postprocess_nodes(\n",
    "            vector_nodes, query_bundle\n",
    "        )\n",
    "\n",
    "        return retrieved_nodes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create Nodes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.node_parser import SentenceSplitter\n",
    "\n",
    "node_parser = SentenceSplitter(chunk_size=1024, chunk_overlap=200)\n",
    "\n",
    "nodes = node_parser.get_nodes_from_documents(documents, show_progress=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set node ids\n",
    "\n",
    "Useful to have consistent result comparison for nodes with and without contextual text."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set node ids\n",
    "\n",
    "nodes = set_node_ids(nodes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'file_path': 'paul_graham_essay.txt',\n",
       " 'file_name': 'paul_graham_essay.txt',\n",
       " 'file_type': 'text/plain',\n",
       " 'file_size': 75042,\n",
       " 'creation_date': '2024-10-01',\n",
       " 'last_modified_date': '2024-10-01'}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nodes[0].metadata"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create contextual nodes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nodes_contextual = create_contextual_nodes(nodes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "({'file_path': 'paul_graham_essay.txt',\n",
       "  'file_name': 'paul_graham_essay.txt',\n",
       "  'file_type': 'text/plain',\n",
       "  'file_size': 75042,\n",
       "  'creation_date': '2024-10-01',\n",
       "  'last_modified_date': '2024-10-01'},\n",
       " {'file_path': 'paul_graham_essay.txt',\n",
       "  'file_name': 'paul_graham_essay.txt',\n",
       "  'file_type': 'text/plain',\n",
       "  'file_size': 75042,\n",
       "  'creation_date': '2024-10-01',\n",
       "  'last_modified_date': '2024-10-01',\n",
       "  'context': 'assistant: This chunk is the opening section of Paul Graham\\'s essay \"What I Worked On,\" describing his early experiences with programming and writing as a teenager, his initial interest in philosophy in college, and his subsequent shift to studying artificial intelligence in the mid-1980s.'})"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nodes[0].metadata, nodes_contextual[0].metadata"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set `similarity_top_k`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "similarity_top_k = 3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set `CohereReranker`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.postprocessor.cohere_rerank import CohereRerank\n",
    "\n",
    "cohere_rerank = CohereRerank(\n",
    "    api_key=os.environ[\"COHERE_API_KEY\"], top_n=similarity_top_k\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create retrievers.\n",
    "\n",
    "1. Embedding based retriever.\n",
    "2. BM25 based retriever.\n",
    "3. Embedding + BM25 + Cohere reranker retriever."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "DEBUG:bm25s:Building index from IDs objects\n"
     ]
    }
   ],
   "source": [
    "embedding_retriever = create_embedding_retriever(\n",
    "    nodes, similarity_top_k=similarity_top_k\n",
    ")\n",
    "bm25_retriever = create_bm25_retriever(\n",
    "    nodes, similarity_top_k=similarity_top_k\n",
    ")\n",
    "embedding_bm25_retriever_rerank = EmbeddingBM25RerankerRetriever(\n",
    "    embedding_retriever, bm25_retriever, reranker=cohere_rerank\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create retrievers using contextual nodes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "DEBUG:bm25s:Building index from IDs objects\n"
     ]
    }
   ],
   "source": [
    "contextual_embedding_retriever = create_embedding_retriever(\n",
    "    nodes_contextual, similarity_top_k=similarity_top_k\n",
    ")\n",
    "contextual_bm25_retriever = create_bm25_retriever(\n",
    "    nodes_contextual, similarity_top_k=similarity_top_k\n",
    ")\n",
    "contextual_embedding_bm25_retriever_rerank = EmbeddingBM25RerankerRetriever(\n",
    "    contextual_embedding_retriever,\n",
    "    contextual_bm25_retriever,\n",
    "    reranker=cohere_rerank,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create Synthetic query dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 21/21 [02:59<00:00,  8.53s/it]\n"
     ]
    }
   ],
   "source": [
    "from llama_index.llms.openai import OpenAI\n",
    "\n",
    "llm = OpenAI(model=\"gpt-4\")\n",
    "\n",
    "qa_dataset = create_eval_dataset(nodes, llm=llm, num_questions_per_chunk=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.google.colaboratory.intrinsic+json": {
       "type": "string"
      },
      "text/plain": [
       "\"The author initially intended to study philosophy in college but later switched to AI. Discuss the reasons behind this shift in interest and how specific influences like Heinlein's novel and Winograd's SHRDLU contributed to his decision.\""
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "list(qa_dataset.queries.values())[1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluate retrievers with and without contextual nodes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "embedding_retriever_results = await retrieval_results(\n",
    "    embedding_retriever, qa_dataset\n",
    ")\n",
    "bm25_retriever_results = await retrieval_results(bm25_retriever, qa_dataset)\n",
    "embedding_bm25_retriever_rerank_results = await retrieval_results(\n",
    "    embedding_bm25_retriever_rerank, qa_dataset\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "contextual_embedding_retriever_results = await retrieval_results(\n",
    "    contextual_embedding_retriever, qa_dataset\n",
    ")\n",
    "contextual_bm25_retriever_results = await retrieval_results(\n",
    "    contextual_bm25_retriever, qa_dataset\n",
    ")\n",
    "contextual_embedding_bm25_retriever_rerank_results = await retrieval_results(\n",
    "    contextual_embedding_bm25_retriever_rerank, qa_dataset\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Display results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Without Context"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.google.colaboratory.intrinsic+json": {
       "summary": "{\n  \"name\": \"pd\",\n  \"rows\": 3,\n  \"fields\": [\n    {\n      \"column\": \"retrievers\",\n      \"properties\": {\n        \"dtype\": \"string\",\n        \"num_unique_values\": 3,\n        \"samples\": [\n          \"Embedding Retriever\",\n          \"BM25 Retriever\",\n          \"Embedding + BM25 Retriever + Reranker\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"hit_rate\",\n      \"properties\": {\n        \"dtype\": \"number\",\n        \"std\": 0.047619047619047616,\n        \"min\": 0.8571428571428571,\n        \"max\": 0.9523809523809523,\n        \"num_unique_values\": 3,\n        \"samples\": [\n          0.8571428571428571,\n          0.9047619047619048,\n          0.9523809523809523\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"mrr\",\n      \"properties\": {\n        \"dtype\": \"number\",\n        \"std\": 0.07020557941648461,\n        \"min\": 0.7261904761904762,\n        \"max\": 0.865079365079365,\n        \"num_unique_values\": 3,\n        \"samples\": [\n          0.7261904761904762,\n          0.7777777777777779,\n          0.865079365079365\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"precision\",\n      \"properties\": {\n        \"dtype\": \"number\",\n        \"std\": 0.09426865504019484,\n        \"min\": 0.2857142857142857,\n        \"max\": 0.4563492063492063,\n        \"num_unique_values\": 3,\n        \"samples\": [\n          0.2857142857142857,\n          0.3015873015873016,\n          0.4563492063492063\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"recall\",\n      \"properties\": {\n        \"dtype\": \"number\",\n        \"std\": 0.047619047619047616,\n        \"min\": 0.8571428571428571,\n        \"max\": 0.9523809523809523,\n        \"num_unique_values\": 3,\n        \"samples\": [\n          0.8571428571428571,\n          0.9047619047619048,\n          0.9523809523809523\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"ap\",\n      \"properties\": {\n        \"dtype\": \"number\",\n        \"std\": 0.07020557941648461,\n        \"min\": 0.7261904761904762,\n        \"max\": 0.865079365079365,\n        \"num_unique_values\": 3,\n        \"samples\": [\n          0.7261904761904762,\n          0.7777777777777779,\n          0.865079365079365\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"ndcg\",\n      \"properties\": {\n        \"dtype\": \"number\",\n        \"std\": 0.09414681939870749,\n        \"min\": 0.3566128176758776,\n        \"max\": 0.5301723159805233,\n        \"num_unique_values\": 3,\n        \"samples\": [\n          0.3566128176758776,\n          0.38015732175191164,\n          0.5301723159805233\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    }\n  ]\n}",
       "type": "dataframe"
      },
      "text/html": [
       "\n",
       "  <div id=\"df-d3120688-09ee-4895-b98e-b26fe4dab8b3\" class=\"colab-df-container\">\n",
       "    <div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>retrievers</th>\n",
       "      <th>hit_rate</th>\n",
       "      <th>mrr</th>\n",
       "      <th>precision</th>\n",
       "      <th>recall</th>\n",
       "      <th>ap</th>\n",
       "      <th>ndcg</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Embedding Retriever</td>\n",
       "      <td>0.857143</td>\n",
       "      <td>0.726190</td>\n",
       "      <td>0.285714</td>\n",
       "      <td>0.857143</td>\n",
       "      <td>0.726190</td>\n",
       "      <td>0.356613</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>BM25 Retriever</td>\n",
       "      <td>0.904762</td>\n",
       "      <td>0.777778</td>\n",
       "      <td>0.301587</td>\n",
       "      <td>0.904762</td>\n",
       "      <td>0.777778</td>\n",
       "      <td>0.380157</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Embedding + BM25 Retriever + Reranker</td>\n",
       "      <td>0.952381</td>\n",
       "      <td>0.865079</td>\n",
       "      <td>0.456349</td>\n",
       "      <td>0.952381</td>\n",
       "      <td>0.865079</td>\n",
       "      <td>0.530172</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>\n",
       "    <div class=\"colab-df-buttons\">\n",
       "\n",
       "  <div class=\"colab-df-container\">\n",
       "    <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-d3120688-09ee-4895-b98e-b26fe4dab8b3')\"\n",
       "            title=\"Convert this dataframe to an interactive table.\"\n",
       "            style=\"display:none;\">\n",
       "\n",
       "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\" viewBox=\"0 -960 960 960\">\n",
       "    <path d=\"M120-120v-720h720v720H120Zm60-500h600v-160H180v160Zm220 220h160v-160H400v160Zm0 220h160v-160H400v160ZM180-400h160v-160H180v160Zm440 0h160v-160H620v160ZM180-180h160v-160H180v160Zm440 0h160v-160H620v160Z\"/>\n",
       "  </svg>\n",
       "    </button>\n",
       "\n",
       "  <style>\n",
       "    .colab-df-container {\n",
       "      display:flex;\n",
       "      gap: 12px;\n",
       "    }\n",
       "\n",
       "    .colab-df-convert {\n",
       "      background-color: #E8F0FE;\n",
       "      border: none;\n",
       "      border-radius: 50%;\n",
       "      cursor: pointer;\n",
       "      display: none;\n",
       "      fill: #1967D2;\n",
       "      height: 32px;\n",
       "      padding: 0 0 0 0;\n",
       "      width: 32px;\n",
       "    }\n",
       "\n",
       "    .colab-df-convert:hover {\n",
       "      background-color: #E2EBFA;\n",
       "      box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
       "      fill: #174EA6;\n",
       "    }\n",
       "\n",
       "    .colab-df-buttons div {\n",
       "      margin-bottom: 4px;\n",
       "    }\n",
       "\n",
       "    [theme=dark] .colab-df-convert {\n",
       "      background-color: #3B4455;\n",
       "      fill: #D2E3FC;\n",
       "    }\n",
       "\n",
       "    [theme=dark] .colab-df-convert:hover {\n",
       "      background-color: #434B5C;\n",
       "      box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
       "      filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
       "      fill: #FFFFFF;\n",
       "    }\n",
       "  </style>\n",
       "\n",
       "    <script>\n",
       "      const buttonEl =\n",
       "        document.querySelector('#df-d3120688-09ee-4895-b98e-b26fe4dab8b3 button.colab-df-convert');\n",
       "      buttonEl.style.display =\n",
       "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
       "\n",
       "      async function convertToInteractive(key) {\n",
       "        const element = document.querySelector('#df-d3120688-09ee-4895-b98e-b26fe4dab8b3');\n",
       "        const dataTable =\n",
       "          await google.colab.kernel.invokeFunction('convertToInteractive',\n",
       "                                                    [key], {});\n",
       "        if (!dataTable) return;\n",
       "\n",
       "        const docLinkHtml = 'Like what you see? Visit the ' +\n",
       "          '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
       "          + ' to learn more about interactive tables.';\n",
       "        element.innerHTML = '';\n",
       "        dataTable['output_type'] = 'display_data';\n",
       "        await google.colab.output.renderOutput(dataTable, element);\n",
       "        const docLink = document.createElement('div');\n",
       "        docLink.innerHTML = docLinkHtml;\n",
       "        element.appendChild(docLink);\n",
       "      }\n",
       "    </script>\n",
       "  </div>\n",
       "\n",
       "\n",
       "<div id=\"df-8a4f961a-c4a0-4da5-a4c3-f482dfe971ce\">\n",
       "  <button class=\"colab-df-quickchart\" onclick=\"quickchart('df-8a4f961a-c4a0-4da5-a4c3-f482dfe971ce')\"\n",
       "            title=\"Suggest charts\"\n",
       "            style=\"display:none;\">\n",
       "\n",
       "<svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
       "     width=\"24px\">\n",
       "    <g>\n",
       "        <path d=\"M19 3H5c-1.1 0-2 .9-2 2v14c0 1.1.9 2 2 2h14c1.1 0 2-.9 2-2V5c0-1.1-.9-2-2-2zM9 17H7v-7h2v7zm4 0h-2V7h2v10zm4 0h-2v-4h2v4z\"/>\n",
       "    </g>\n",
       "</svg>\n",
       "  </button>\n",
       "\n",
       "<style>\n",
       "  .colab-df-quickchart {\n",
       "      --bg-color: #E8F0FE;\n",
       "      --fill-color: #1967D2;\n",
       "      --hover-bg-color: #E2EBFA;\n",
       "      --hover-fill-color: #174EA6;\n",
       "      --disabled-fill-color: #AAA;\n",
       "      --disabled-bg-color: #DDD;\n",
       "  }\n",
       "\n",
       "  [theme=dark] .colab-df-quickchart {\n",
       "      --bg-color: #3B4455;\n",
       "      --fill-color: #D2E3FC;\n",
       "      --hover-bg-color: #434B5C;\n",
       "      --hover-fill-color: #FFFFFF;\n",
       "      --disabled-bg-color: #3B4455;\n",
       "      --disabled-fill-color: #666;\n",
       "  }\n",
       "\n",
       "  .colab-df-quickchart {\n",
       "    background-color: var(--bg-color);\n",
       "    border: none;\n",
       "    border-radius: 50%;\n",
       "    cursor: pointer;\n",
       "    display: none;\n",
       "    fill: var(--fill-color);\n",
       "    height: 32px;\n",
       "    padding: 0;\n",
       "    width: 32px;\n",
       "  }\n",
       "\n",
       "  .colab-df-quickchart:hover {\n",
       "    background-color: var(--hover-bg-color);\n",
       "    box-shadow: 0 1px 2px rgba(60, 64, 67, 0.3), 0 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
       "    fill: var(--button-hover-fill-color);\n",
       "  }\n",
       "\n",
       "  .colab-df-quickchart-complete:disabled,\n",
       "  .colab-df-quickchart-complete:disabled:hover {\n",
       "    background-color: var(--disabled-bg-color);\n",
       "    fill: var(--disabled-fill-color);\n",
       "    box-shadow: none;\n",
       "  }\n",
       "\n",
       "  .colab-df-spinner {\n",
       "    border: 2px solid var(--fill-color);\n",
       "    border-color: transparent;\n",
       "    border-bottom-color: var(--fill-color);\n",
       "    animation:\n",
       "      spin 1s steps(1) infinite;\n",
       "  }\n",
       "\n",
       "  @keyframes spin {\n",
       "    0% {\n",
       "      border-color: transparent;\n",
       "      border-bottom-color: var(--fill-color);\n",
       "      border-left-color: var(--fill-color);\n",
       "    }\n",
       "    20% {\n",
       "      border-color: transparent;\n",
       "      border-left-color: var(--fill-color);\n",
       "      border-top-color: var(--fill-color);\n",
       "    }\n",
       "    30% {\n",
       "      border-color: transparent;\n",
       "      border-left-color: var(--fill-color);\n",
       "      border-top-color: var(--fill-color);\n",
       "      border-right-color: var(--fill-color);\n",
       "    }\n",
       "    40% {\n",
       "      border-color: transparent;\n",
       "      border-right-color: var(--fill-color);\n",
       "      border-top-color: var(--fill-color);\n",
       "    }\n",
       "    60% {\n",
       "      border-color: transparent;\n",
       "      border-right-color: var(--fill-color);\n",
       "    }\n",
       "    80% {\n",
       "      border-color: transparent;\n",
       "      border-right-color: var(--fill-color);\n",
       "      border-bottom-color: var(--fill-color);\n",
       "    }\n",
       "    90% {\n",
       "      border-color: transparent;\n",
       "      border-bottom-color: var(--fill-color);\n",
       "    }\n",
       "  }\n",
       "</style>\n",
       "\n",
       "  <script>\n",
       "    async function quickchart(key) {\n",
       "      const quickchartButtonEl =\n",
       "        document.querySelector('#' + key + ' button');\n",
       "      quickchartButtonEl.disabled = true;  // To prevent multiple clicks.\n",
       "      quickchartButtonEl.classList.add('colab-df-spinner');\n",
       "      try {\n",
       "        const charts = await google.colab.kernel.invokeFunction(\n",
       "            'suggestCharts', [key], {});\n",
       "      } catch (error) {\n",
       "        console.error('Error during call to suggestCharts:', error);\n",
       "      }\n",
       "      quickchartButtonEl.classList.remove('colab-df-spinner');\n",
       "      quickchartButtonEl.classList.add('colab-df-quickchart-complete');\n",
       "    }\n",
       "    (() => {\n",
       "      let quickchartButtonEl =\n",
       "        document.querySelector('#df-8a4f961a-c4a0-4da5-a4c3-f482dfe971ce button');\n",
       "      quickchartButtonEl.style.display =\n",
       "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
       "    })();\n",
       "  </script>\n",
       "</div>\n",
       "\n",
       "    </div>\n",
       "  </div>\n"
      ],
      "text/plain": [
       "                              retrievers  hit_rate       mrr  precision  \\\n",
       "0                    Embedding Retriever  0.857143  0.726190   0.285714   \n",
       "1                         BM25 Retriever  0.904762  0.777778   0.301587   \n",
       "2  Embedding + BM25 Retriever + Reranker  0.952381  0.865079   0.456349   \n",
       "\n",
       "     recall        ap      ndcg  \n",
       "0  0.857143  0.726190  0.356613  \n",
       "1  0.904762  0.777778  0.380157  \n",
       "2  0.952381  0.865079  0.530172  "
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pd.concat(\n",
    "    [\n",
    "        display_results(\"Embedding Retriever\", embedding_retriever_results),\n",
    "        display_results(\"BM25 Retriever\", bm25_retriever_results),\n",
    "        display_results(\n",
    "            \"Embedding + BM25 Retriever + Reranker\",\n",
    "            embedding_bm25_retriever_rerank_results,\n",
    "        ),\n",
    "    ],\n",
    "    ignore_index=True,\n",
    "    axis=0,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### With Context"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.google.colaboratory.intrinsic+json": {
       "summary": "{\n  \"name\": \"pd\",\n  \"rows\": 3,\n  \"fields\": [\n    {\n      \"column\": \"retrievers\",\n      \"properties\": {\n        \"dtype\": \"string\",\n        \"num_unique_values\": 3,\n        \"samples\": [\n          \"Contextual Embedding Retriever\",\n          \"Contextual BM25 Retriever\",\n          \"Contextual Embedding + Contextual BM25 Retriever + Reranker\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"hit_rate\",\n      \"properties\": {\n        \"dtype\": \"number\",\n        \"std\": 0.02380952380952378,\n        \"min\": 0.9285714285714286,\n        \"max\": 0.9761904761904762,\n        \"num_unique_values\": 3,\n        \"samples\": [\n          0.9285714285714286,\n          0.9523809523809523,\n          0.9761904761904762\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"mrr\",\n      \"properties\": {\n        \"dtype\": \"number\",\n        \"std\": 0.07745722736485365,\n        \"min\": 0.746031746031746,\n        \"max\": 0.9007936507936507,\n        \"num_unique_values\": 3,\n        \"samples\": [\n          0.746031746031746,\n          0.8293650793650794,\n          0.9007936507936507\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"precision\",\n      \"properties\": {\n        \"dtype\": \"number\",\n        \"std\": 0.09401775472971882,\n        \"min\": 0.30952380952380953,\n        \"max\": 0.47619047619047616,\n        \"num_unique_values\": 3,\n        \"samples\": [\n          0.30952380952380953,\n          0.3174603174603175,\n          0.47619047619047616\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"recall\",\n      \"properties\": {\n        \"dtype\": \"number\",\n        \"std\": 0.02380952380952378,\n        \"min\": 0.9285714285714286,\n        \"max\": 0.9761904761904762,\n        \"num_unique_values\": 3,\n        \"samples\": [\n          0.9285714285714286,\n          0.9523809523809523,\n          0.9761904761904762\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"ap\",\n      \"properties\": {\n        \"dtype\": \"number\",\n        \"std\": 0.07745722736485365,\n        \"min\": 0.746031746031746,\n        \"max\": 0.9007936507936507,\n        \"num_unique_values\": 3,\n        \"samples\": [\n          0.746031746031746,\n          0.8293650793650794,\n          0.9007936507936507\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"ndcg\",\n      \"properties\": {\n        \"dtype\": \"number\",\n        \"std\": 0.09810349914041294,\n        \"min\": 0.37217487410173994,\n        \"max\": 0.555745671725946,\n        \"num_unique_values\": 3,\n        \"samples\": [\n          0.37217487410173994,\n          0.40396684556143553,\n          0.555745671725946\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    }\n  ]\n}",
       "type": "dataframe"
      },
      "text/html": [
       "\n",
       "  <div id=\"df-9344898e-8549-43d4-9f09-e83f6834b765\" class=\"colab-df-container\">\n",
       "    <div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>retrievers</th>\n",
       "      <th>hit_rate</th>\n",
       "      <th>mrr</th>\n",
       "      <th>precision</th>\n",
       "      <th>recall</th>\n",
       "      <th>ap</th>\n",
       "      <th>ndcg</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Contextual Embedding Retriever</td>\n",
       "      <td>0.928571</td>\n",
       "      <td>0.746032</td>\n",
       "      <td>0.309524</td>\n",
       "      <td>0.928571</td>\n",
       "      <td>0.746032</td>\n",
       "      <td>0.372175</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Contextual BM25 Retriever</td>\n",
       "      <td>0.952381</td>\n",
       "      <td>0.829365</td>\n",
       "      <td>0.317460</td>\n",
       "      <td>0.952381</td>\n",
       "      <td>0.829365</td>\n",
       "      <td>0.403967</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Contextual Embedding + Contextual BM25 Retriev...</td>\n",
       "      <td>0.976190</td>\n",
       "      <td>0.900794</td>\n",
       "      <td>0.476190</td>\n",
       "      <td>0.976190</td>\n",
       "      <td>0.900794</td>\n",
       "      <td>0.555746</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>\n",
       "    <div class=\"colab-df-buttons\">\n",
       "\n",
       "  <div class=\"colab-df-container\">\n",
       "    <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-9344898e-8549-43d4-9f09-e83f6834b765')\"\n",
       "            title=\"Convert this dataframe to an interactive table.\"\n",
       "            style=\"display:none;\">\n",
       "\n",
       "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\" viewBox=\"0 -960 960 960\">\n",
       "    <path d=\"M120-120v-720h720v720H120Zm60-500h600v-160H180v160Zm220 220h160v-160H400v160Zm0 220h160v-160H400v160ZM180-400h160v-160H180v160Zm440 0h160v-160H620v160ZM180-180h160v-160H180v160Zm440 0h160v-160H620v160Z\"/>\n",
       "  </svg>\n",
       "    </button>\n",
       "\n",
       "  <style>\n",
       "    .colab-df-container {\n",
       "      display:flex;\n",
       "      gap: 12px;\n",
       "    }\n",
       "\n",
       "    .colab-df-convert {\n",
       "      background-color: #E8F0FE;\n",
       "      border: none;\n",
       "      border-radius: 50%;\n",
       "      cursor: pointer;\n",
       "      display: none;\n",
       "      fill: #1967D2;\n",
       "      height: 32px;\n",
       "      padding: 0 0 0 0;\n",
       "      width: 32px;\n",
       "    }\n",
       "\n",
       "    .colab-df-convert:hover {\n",
       "      background-color: #E2EBFA;\n",
       "      box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
       "      fill: #174EA6;\n",
       "    }\n",
       "\n",
       "    .colab-df-buttons div {\n",
       "      margin-bottom: 4px;\n",
       "    }\n",
       "\n",
       "    [theme=dark] .colab-df-convert {\n",
       "      background-color: #3B4455;\n",
       "      fill: #D2E3FC;\n",
       "    }\n",
       "\n",
       "    [theme=dark] .colab-df-convert:hover {\n",
       "      background-color: #434B5C;\n",
       "      box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
       "      filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
       "      fill: #FFFFFF;\n",
       "    }\n",
       "  </style>\n",
       "\n",
       "    <script>\n",
       "      const buttonEl =\n",
       "        document.querySelector('#df-9344898e-8549-43d4-9f09-e83f6834b765 button.colab-df-convert');\n",
       "      buttonEl.style.display =\n",
       "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
       "\n",
       "      async function convertToInteractive(key) {\n",
       "        const element = document.querySelector('#df-9344898e-8549-43d4-9f09-e83f6834b765');\n",
       "        const dataTable =\n",
       "          await google.colab.kernel.invokeFunction('convertToInteractive',\n",
       "                                                    [key], {});\n",
       "        if (!dataTable) return;\n",
       "\n",
       "        const docLinkHtml = 'Like what you see? Visit the ' +\n",
       "          '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
       "          + ' to learn more about interactive tables.';\n",
       "        element.innerHTML = '';\n",
       "        dataTable['output_type'] = 'display_data';\n",
       "        await google.colab.output.renderOutput(dataTable, element);\n",
       "        const docLink = document.createElement('div');\n",
       "        docLink.innerHTML = docLinkHtml;\n",
       "        element.appendChild(docLink);\n",
       "      }\n",
       "    </script>\n",
       "  </div>\n",
       "\n",
       "\n",
       "<div id=\"df-e4cc9f82-a5c7-4e01-87a4-e2a55d52b61c\">\n",
       "  <button class=\"colab-df-quickchart\" onclick=\"quickchart('df-e4cc9f82-a5c7-4e01-87a4-e2a55d52b61c')\"\n",
       "            title=\"Suggest charts\"\n",
       "            style=\"display:none;\">\n",
       "\n",
       "<svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
       "     width=\"24px\">\n",
       "    <g>\n",
       "        <path d=\"M19 3H5c-1.1 0-2 .9-2 2v14c0 1.1.9 2 2 2h14c1.1 0 2-.9 2-2V5c0-1.1-.9-2-2-2zM9 17H7v-7h2v7zm4 0h-2V7h2v10zm4 0h-2v-4h2v4z\"/>\n",
       "    </g>\n",
       "</svg>\n",
       "  </button>\n",
       "\n",
       "<style>\n",
       "  .colab-df-quickchart {\n",
       "      --bg-color: #E8F0FE;\n",
       "      --fill-color: #1967D2;\n",
       "      --hover-bg-color: #E2EBFA;\n",
       "      --hover-fill-color: #174EA6;\n",
       "      --disabled-fill-color: #AAA;\n",
       "      --disabled-bg-color: #DDD;\n",
       "  }\n",
       "\n",
       "  [theme=dark] .colab-df-quickchart {\n",
       "      --bg-color: #3B4455;\n",
       "      --fill-color: #D2E3FC;\n",
       "      --hover-bg-color: #434B5C;\n",
       "      --hover-fill-color: #FFFFFF;\n",
       "      --disabled-bg-color: #3B4455;\n",
       "      --disabled-fill-color: #666;\n",
       "  }\n",
       "\n",
       "  .colab-df-quickchart {\n",
       "    background-color: var(--bg-color);\n",
       "    border: none;\n",
       "    border-radius: 50%;\n",
       "    cursor: pointer;\n",
       "    display: none;\n",
       "    fill: var(--fill-color);\n",
       "    height: 32px;\n",
       "    padding: 0;\n",
       "    width: 32px;\n",
       "  }\n",
       "\n",
       "  .colab-df-quickchart:hover {\n",
       "    background-color: var(--hover-bg-color);\n",
       "    box-shadow: 0 1px 2px rgba(60, 64, 67, 0.3), 0 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
       "    fill: var(--button-hover-fill-color);\n",
       "  }\n",
       "\n",
       "  .colab-df-quickchart-complete:disabled,\n",
       "  .colab-df-quickchart-complete:disabled:hover {\n",
       "    background-color: var(--disabled-bg-color);\n",
       "    fill: var(--disabled-fill-color);\n",
       "    box-shadow: none;\n",
       "  }\n",
       "\n",
       "  .colab-df-spinner {\n",
       "    border: 2px solid var(--fill-color);\n",
       "    border-color: transparent;\n",
       "    border-bottom-color: var(--fill-color);\n",
       "    animation:\n",
       "      spin 1s steps(1) infinite;\n",
       "  }\n",
       "\n",
       "  @keyframes spin {\n",
       "    0% {\n",
       "      border-color: transparent;\n",
       "      border-bottom-color: var(--fill-color);\n",
       "      border-left-color: var(--fill-color);\n",
       "    }\n",
       "    20% {\n",
       "      border-color: transparent;\n",
       "      border-left-color: var(--fill-color);\n",
       "      border-top-color: var(--fill-color);\n",
       "    }\n",
       "    30% {\n",
       "      border-color: transparent;\n",
       "      border-left-color: var(--fill-color);\n",
       "      border-top-color: var(--fill-color);\n",
       "      border-right-color: var(--fill-color);\n",
       "    }\n",
       "    40% {\n",
       "      border-color: transparent;\n",
       "      border-right-color: var(--fill-color);\n",
       "      border-top-color: var(--fill-color);\n",
       "    }\n",
       "    60% {\n",
       "      border-color: transparent;\n",
       "      border-right-color: var(--fill-color);\n",
       "    }\n",
       "    80% {\n",
       "      border-color: transparent;\n",
       "      border-right-color: var(--fill-color);\n",
       "      border-bottom-color: var(--fill-color);\n",
       "    }\n",
       "    90% {\n",
       "      border-color: transparent;\n",
       "      border-bottom-color: var(--fill-color);\n",
       "    }\n",
       "  }\n",
       "</style>\n",
       "\n",
       "  <script>\n",
       "    async function quickchart(key) {\n",
       "      const quickchartButtonEl =\n",
       "        document.querySelector('#' + key + ' button');\n",
       "      quickchartButtonEl.disabled = true;  // To prevent multiple clicks.\n",
       "      quickchartButtonEl.classList.add('colab-df-spinner');\n",
       "      try {\n",
       "        const charts = await google.colab.kernel.invokeFunction(\n",
       "            'suggestCharts', [key], {});\n",
       "      } catch (error) {\n",
       "        console.error('Error during call to suggestCharts:', error);\n",
       "      }\n",
       "      quickchartButtonEl.classList.remove('colab-df-spinner');\n",
       "      quickchartButtonEl.classList.add('colab-df-quickchart-complete');\n",
       "    }\n",
       "    (() => {\n",
       "      let quickchartButtonEl =\n",
       "        document.querySelector('#df-e4cc9f82-a5c7-4e01-87a4-e2a55d52b61c button');\n",
       "      quickchartButtonEl.style.display =\n",
       "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
       "    })();\n",
       "  </script>\n",
       "</div>\n",
       "\n",
       "    </div>\n",
       "  </div>\n"
      ],
      "text/plain": [
       "                                          retrievers  hit_rate       mrr  \\\n",
       "0                     Contextual Embedding Retriever  0.928571  0.746032   \n",
       "1                          Contextual BM25 Retriever  0.952381  0.829365   \n",
       "2  Contextual Embedding + Contextual BM25 Retriev...  0.976190  0.900794   \n",
       "\n",
       "   precision    recall        ap      ndcg  \n",
       "0   0.309524  0.928571  0.746032  0.372175  \n",
       "1   0.317460  0.952381  0.829365  0.403967  \n",
       "2   0.476190  0.976190  0.900794  0.555746  "
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pd.concat(\n",
    "    [\n",
    "        display_results(\n",
    "            \"Contextual Embedding Retriever\",\n",
    "            contextual_embedding_retriever_results,\n",
    "        ),\n",
    "        display_results(\n",
    "            \"Contextual BM25 Retriever\", contextual_bm25_retriever_results\n",
    "        ),\n",
    "        display_results(\n",
    "            \"Contextual Embedding + Contextual BM25 Retriever + Reranker\",\n",
    "            contextual_embedding_bm25_retriever_rerank_results,\n",
    "        ),\n",
    "    ],\n",
    "    ignore_index=True,\n",
    "    axis=0,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Observation:\n",
    "\n",
    "We observed improved metrics with contextual retrieval; however, our experiments showed that much depends on the queries, chunk size, chunk overlap, and other variables. Therefore, it’s essential to experiment to optimize the benefits of this technique."
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "llama-index-caVs7DDe-py3.10",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
